from utils.functions import *

class YOLOV5Infer:
    def __init__(self,infer_size,conf_thres, iou_thres, classes, agnostic_nms, max_det,device):
        self.infer_size = infer_size
        self.model = attempt_load("./weights/yolov5s.pt",device,fuse=True)
        self.conf_thresh = conf_thres
        self.iou_thresh = iou_thres
        self.classes = classes
        self.agnostic_nms = agnostic_nms
        self.max_det = max_det

    def prepocess(self,image):
        h,w,c = image.shape
        scale = min(float(self.infer_size)/w,float(self.infer_size)/h)
        w1 = int(w * scale)
        h1 = int(h * scale)
        image_resize = cv2.resize(image,(w1,h1),cv2.INTER_LINEAR)
        # bgr2rgb
        image_bgr = image_resize[:,:,::-1]
        # copy and paste
        image_pad = np.ones((self.infer_size,self.infer_size,3),dtype=np.uint8)*127
        image_pad[(self.infer_size-h1)//2:(self.infer_size-h1)//2 + h1,(self.infer_size-w1)//2:(self.infer_size-w1)//2 + w1,:] = image_bgr
        # transpose  HWC2CHW
        image_trans = np.transpose(image_pad,(2,0,1))
        # CHW-->[1,C,H,W]
        image_batch = np.expand_dims(image_trans,axis=0)
        # normlize
        image_norm = image_batch / 255.0
        # to gpu input tensor
        input_tensor = torch.from_numpy(image_norm).to(device).to(torch.float32)

        return input_tensor
    
    
    
    def postprepocess(self,out,infer_size,orin_size):
        # memcpy from device to host
        out.clone().detach()
        # nms
        pred = non_max_suppression(out, self.conf_thresh, self.iou_thresh, self.classes, self.agnostic_nms, self.max_det)
        # coord_trans
        res = []
        for i,det in enumerate(pred):
            if len(det):
                det[:,:4] = scale_coords(infer_size,det[:,:4],orin_size).round()
                res.append(det.detach().cpu().numpy())
        
        return res
    

    def __call__(self,x):
        input_tensor = self.prepocess(x)
        out = self.model(input_tensor)[0]
        res = self.postprepocess(out,(self.infer_size,self.infer_size),x.shape)
        return res






if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    img_dir = "images/bus.jpg"
    image0 = cv2.imread(img_dir)
    device_ = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    infer_size_ = 640
    conf_thresh_ = 0.25
    iou_thresh_ = 0.65
    classes_ = None
    agnostic_nms_ = False
    max_det_ = 1000
    yolov5 = YOLOV5Infer(infer_size_,conf_thresh_,iou_thresh_,classes_,agnostic_nms_,max_det_,device_)
    res = yolov5(image0)[0]
    color = Colors()
    img_plot = plot_image(res,image0,color)
    cv2.imwrite("img_plot.png",img_plot)